from sklearn.preprocessing import LabelEncoder, OneHotEncoder, StandardScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
import pandas as pd

from Config import Config


def label_encode(x, cols):
	for i in cols:
		labelencoder = LabelEncoder()
		x[:, i] = labelencoder.fit_transform(x[:, i])

def one_hot_encode(x, cols):
	for i in cols:
		onehotencoder = OneHotEncoder(categorical_features = [i])
		x = onehotencoder.fit_transform(x).toarray()

def get_data():

    print("Loading data...")
    data = pd.read_csv(Config().DATA_PATH)

    # 二分类
    if(Config.if_multi == False):
        data['normal.'] = data['normal.'].replace(['back.', 'buffer_overflow.', 'ftp_write.', 'guess_passwd.', 'imap.', 'ipsweep.', 'land.', 'loadmodule.', 'multihop.', 'neptune.', 'nmap.', 'perl.', 'phf.', 'pod.', 'portsweep.', 'rootkit.', 'satan.', 'smurf.', 'spy.', 'teardrop.', 'warezclient.', 'warezmaster.'], 'attack')
    # 四分类
    else:
        data['normal.'] = data['normal.'].replace(['back.', 'land.', 'neptune.', 'pod.', 'smurf.', 'teardrop.'], 'dos').replace(['buffer_overflow.', 'loadmodule.', 'perl.', 'rootkit.'], 'u2r').replace(['ftp_write.', 'guess_passwd.', 'imap.', 'multihop.', 'phf.', 'spy.', 'warezclient.', 'warezmaster'], 'r2l').replace(['ipsweep.', 'nmap.', 'portsweep.', 'satan.'], 'probe')

    X = data.iloc[:, :-1].values
    Y = data.iloc[:, 41].values

    # 编码标签
    cols = [1, 2, 3]
    label_encode(X, cols)
    one_hot_encode(X, cols)
    
    labelencoder_y = LabelEncoder()
    Y = labelencoder_y.fit_transform(Y)

    # 标准化数据
    X = StandardScaler().fit_transform(X)
    # PCA 降维
    X = PCA(n_components = 8).fit_transform(X)

    # 划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size = 0.3, random_state = 42)
    return x_train, x_test, y_train, y_test